﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Management;
using System.Runtime.InteropServices;
using System.Text;
using System.Threading;
using System.Windows;
using System.Windows.Interop;
using USBClassLibrary;

namespace Ornithopter.FlashDiskHelper
{
    /// <summary>
    /// 表示一个USB存储设备（硬件设备），该设备可以包含多个逻辑分区（LogicalDisks）
    /// </summary>
    public partial class USBStorage : IDisposable
    {
        /// <summary>
        /// 专门用来通过USBClass.USBDeviceEventArgs来获取USB存储设备的进程
        /// </summary>
        private static Thread _thread = new Thread(new ParameterizedThreadStart((o) =>
        {
            GetUSBDevice((USBClass.USBDeviceEventArgs)o);
        }));

        private static Object _lockObj = new Object();

        /// <summary>
        /// 接入的USB存储设备列表
        /// </summary>
        private static List<USBStorage> _attachedUSBStorages;

        /// <summary>
        /// 处理消息列队的窗口的句柄
        /// </summary>
        private static HwndSource _hwndSource;

        /// <summary>
        /// USB接口
        /// </summary>
        private static USBClass _USBPort;

        private USBClass.DeviceProperties _deviceProperties;

        /// <summary>
        /// win32_diskdrive，WMI中的磁盘类
        /// </summary>
        private ManagementObject mDiskDrive;

        /// <summary>
        /// 逻辑分区
        /// </summary>
        private List<LogicalDisk> _logicalDisks;

        public String Id { get; private set; }

        /// <summary>
        /// 该设备中的逻辑分区
        /// </summary>
        public List<LogicalDisk> LogicalDisks
        {
            get { return _logicalDisks; }
        }

        /// <summary>
        /// 用Win32_Diskdrive类来初始化一个USBStorage
        /// </summary>
        /// <param name="win32_diskdrive">Win32_Diskdrive的实例</param>
        private USBStorage(ManagementObject win32_diskdrive)
        {
            mDiskDrive = win32_diskdrive;
        }

        /// <summary>
        /// USB接口事件，设备拔出
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="e"></param>
        static void USBDeviceRemoved(object sender, USBClass.USBDeviceEventArgs e)
        {
            lock (_lockObj)
            {
                String info = String.Concat(e.DeviceInterface.dbcc_name);
                if (info.StartsWith("\\\\?\\USBSTOR#"))
                {
                    String id = info.Split('#')[2].Replace("&0", "");
                    USBStorage detached = _attachedUSBStorages.FirstOrDefault(usb => usb.Id == id);
                    if (USBStorageRemoved != null && detached != null)
                    {
                        _attachedUSBStorages.Remove(detached);
                        USBStorageRemoved(detached);

                        detached.Dispose();
                        detached = null;
                        GC.Collect();
                    }
                }
            }
        }

        /// <summary>
        /// USB接口事件，设备请求弹出
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="e"></param>
        static void USBDeviceQueryRemove(object sender, USBClass.USBDeviceEventArgs e)
        {
        }

        /// <summary>
        /// USB接口事件，设备插入并准备就绪
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="e"></param>
        static void USBDeviceAttached(object sender, USBClass.USBDeviceEventArgs e)
        {
            // 另开进程，防止跨线程错误
            new Thread(new ThreadStart(() =>
            {
                GetUSBDevice((USBClass.USBDeviceEventArgs)e);
            })).Start();
        }


        /// <summary>
        /// 从USBClass.USBDeviceEventArgs中获取USB存储设备
        /// </summary>
        /// <param name="e"></param>
        private static void GetUSBDevice(USBClass.USBDeviceEventArgs e)
        {
            lock (_lockObj)
            {
                ManagementObject mObject = null;
                USBClass.DeviceProperties properties = new USBClass.DeviceProperties();
                String info = String.Concat(e.DeviceInterface.dbcc_name);
                String id = null;
                if (info.StartsWith("\\\\?\\USB#") && USBClass.GetUSBDevice(e.DeviceInterface, ref properties, false))
                {
                    id = info.Split('#')[2];

                    // 这里似乎有一个BUG：如果不加上Thread.Sleep(100)，那么在拔掉USB设备再立刻插上
                    // ，程序无法找到该USB设备
                    Thread.Sleep(100);

                    using (ManagementObjectSearcher searcher = new ManagementObjectSearcher("SELECT * FROM Win32_DiskDrive WHERE InterfaceType='USB'"))
                    {
                        try
                        {
                            // 不能再USBClass.USBDeviceAttached事件中直接调用searcher.Get()
                            // 否则会出现跨线程错误
                            using (ManagementObjectCollection moc = searcher.Get())
                            {
                                String pnpDevId = null;
                                
                                foreach (ManagementObject mo in moc)
                                {
                                    pnpDevId = mo.Properties["PNPDeviceID"].Value.ToString();
                                    if (pnpDevId.Split('\\').Last().Replace("&0", "") == id)
                                    {
                                        mObject = mo;
                                    }
                                }
                            }
                        }
                        catch (Exception err)
                        {
                            System.Diagnostics.Debug.WriteLine("An error occured when retrieveing infomation : \n" + err.Message);
                        }
                    }
                    if (mObject != null)
                    {
                        USBStorage usb = new USBStorage(mObject);
                        usb._deviceProperties = properties;
                        usb.Id = id;
                        usb.GetPartitions();
                        if (_attachedUSBStorages == null)
                            _attachedUSBStorages = new List<USBStorage>();
                        _attachedUSBStorages.Add(usb);
                        if (USBStorageAttached != null)
                        {
                            USBStorageAttached(usb);
                        }
                    }
                }
            }
        }

        /// <summary>
        /// 请求弹出设备
        /// </summary>
        public void Eject()
        {
            CM_Request_Device_Eject_NoUi(_deviceProperties.DevInst, IntPtr.Zero, null, 0, 0);
        }

        /// <summary>
        /// 获取逻辑分区
        /// </summary>
        private void GetPartitions()
        {
            this._logicalDisks = new List<LogicalDisk>();
            foreach (ManagementObject part in mDiskDrive.GetRelated("Win32_DiskPartition"))
            {
                foreach (ManagementObject mapping in part.GetRelated("Win32_LogicalDisk"))
                {
                    LogicalDisk disk = new LogicalDisk(mapping, this);
                    this._logicalDisks.Add(disk);
                }
            }
        }

        public static event Action<USBStorage> USBStorageAttached;

        public static event Action<USBStorage> USBStorageRemoved;

        #region Hook
        public static void StartListenning()
        {
            if (Application.Current.MainWindow != null)
            {
                _hwndSource = PresentationSource.FromVisual(Application.Current.MainWindow) as HwndSource;

                if (_hwndSource != null)
                {
                    _hwndSource.AddHook(WndProc);
                }
            }
            if (_USBPort == null && _hwndSource != null)
            {
                _USBPort = new USBClass();
                _USBPort.USBDeviceAttached += USBDeviceAttached;
                _USBPort.USBDeviceQueryRemove += USBDeviceQueryRemove;
                _USBPort.USBDeviceRemoved += USBDeviceRemoved;
                _USBPort.RegisterForDeviceChange(true, _hwndSource.Handle);
            }
        }

        private static IntPtr WndProc(IntPtr hwnd, int msg, IntPtr wParam, IntPtr lParam, ref bool handled)
        {
            handled = false;
            if (_USBPort != null)
            {
                System.Windows.Forms.Message m = new System.Windows.Forms.Message()
                {
                    Msg = msg,
                    WParam = wParam,
                    LParam = lParam
                };
                _USBPort.ProcessWindowsMessage(ref m);
            }
            return IntPtr.Zero;
        }

        public static void StopListenning()
        {
            if (_hwndSource != null)
            {
                if(!_hwndSource.IsDisposed)
                    _hwndSource.Dispose();
                _hwndSource = null;
            }
            if (_USBPort != null)
            {
                _USBPort.USBDeviceAttached -= USBDeviceAttached;
                _USBPort.USBDeviceQueryRemove -= USBDeviceQueryRemove;
                _USBPort.USBDeviceRemoved -= USBDeviceRemoved;
                _USBPort = null;
            }
        }
        #endregion

        #region Native
        // from cfg.h
        internal enum PNP_VETO_TYPE
        {
            Ok,

            TypeUnknown,
            LegacyDevice,
            PendingClose,
            WindowsApp,
            WindowsService,
            OutstandingOpen,
            Device,
            Driver,
            IllegalDeviceRequest,
            InsufficientPower,
            NonDisableable,
            LegacyDriver,
        }

        [DllImport("setupapi.dll")]
        internal static extern int CM_Request_Device_Eject(
            uint dnDevInst,
            out PNP_VETO_TYPE pVetoType,
            StringBuilder pszVetoName,
            int ulNameLength,
            int ulFlags
            );

        [DllImport("setupapi.dll", EntryPoint = "CM_Request_Device_Eject")]
        internal static extern int CM_Request_Device_Eject_NoUi(
            uint dnDevInst,
            IntPtr pVetoType,
            StringBuilder pszVetoName,
            int ulNameLength,
            int ulFlags
            );
        #endregion

        public void Dispose()
        {
            if (mDiskDrive != null)
                mDiskDrive.Dispose();
            if (_logicalDisks != null && _logicalDisks.Count >0)
            {
                _logicalDisks.ForEach(disk => disk.Dispose());
            }
        }
    }
}
